package cn.enneagon.kudu;

import org.apache.kudu.ColumnSchema;
import org.apache.kudu.Schema;
import org.apache.kudu.Type;
import org.apache.kudu.client.*;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.util.ArrayList;

/**
 * 直接使用kudu client的java API的demo
 */

public class testKudu {

    //声明全局变量 方便后面增删改查操作
    private KuduClient kuduClient;

    //kudu master
    private String kuduMaster;

    //kudu中的表名
    private String kuduTable;

    @Before
    public void init() {
        kuduMaster = "cdh001:7051,cdh002:7051,cdh003:7051";

        //指定待操作的table名字
        kuduTable = "Student";

        KuduClient.KuduClientBuilder kuduClientBuilder = new KuduClient.KuduClientBuilder(kuduMaster);
        //指定客户端和kudu集群socket超时时间
        kuduClientBuilder.defaultSocketReadTimeoutMs(10000);

        //通过builder中build方法创建kuduclient
        kuduClient = kuduClientBuilder.build();
    }

    @After
    public void close() throws KuduException {
        if (kuduClient != null) {
            kuduClient.close();
        }
    }

    /**
     * 创建表
     */
    @Test
    public void createTable() throws KuduException {
        //判断表是否存在，不存在就构建
        if (!kuduClient.tableExists(kuduTable)) {

            //构建创建表的schema信息------就是表的字段和类型, id为表的主键
            ArrayList<ColumnSchema> columnSchemas = new ArrayList<>();
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("id", Type.INT32).key(true).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("name", Type.STRING).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("age", Type.INT32).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("sex", Type.INT32).build());
            Schema schema = new Schema(columnSchemas);

            //指定创建表的option属性
            CreateTableOptions tableOptions = new CreateTableOptions();

            //指定表的分区规则 采用hash分区 根据id哈希到指定的6个部分中
            ArrayList<String> partitionList = new ArrayList<>();
            partitionList.add("id");
            tableOptions.addHashPartitions(partitionList, 6); //id.hashcode % 6

            //如果表不存在，进行创建表的操作(表名 表schema信息 表属性信息)
            kuduClient.createTable(kuduTable, schema, tableOptions);
        }
    }

    /**
     * 向表插入数据
     */
    @Test
    public void insertTable() throws KuduException {
        //向表加载数据需要一个 kuduSession 对象
        KuduSession kuduSession = kuduClient.newSession();
        kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_SYNC);

        KuduTable kuduTable = kuduClient.openTable(this.kuduTable);
        for (int i = 0; i < 10; i++) {
            Insert insert = kuduTable.newInsert();
            PartialRow row = insert.getRow();
            row.addInt("id", i);
            row.addString("name", "zhangsan-" + i);
            row.addInt("age", 20+i);
            row.addInt("sex", i%2);
            kuduSession.apply(insert);
        }
    }

    /**
     * 查询表的数据
     */
    @Test
    public void queryData() throws KuduException {
        KuduScanner.KuduScannerBuilder kuduScannerBuilder = kuduClient.newScannerBuilder(kuduClient.openTable(kuduTable));
        ArrayList<String> columnList = new ArrayList<>();
        columnList.add("id");
        columnList.add("name");
        columnList.add("age");
        columnList.add("sex");
        kuduScannerBuilder.setProjectedColumnNames(columnList);

        KuduScanner kuduScanner = kuduScannerBuilder.build();

        while (kuduScanner.hasMoreRows()) {
            RowResultIterator rowResults = kuduScanner.nextRows();

            while (rowResults.hasNext()) {
                RowResult row = rowResults.next();
                int id = row.getInt("id");
                String name = row.getString("name");
                int age = row.getInt("age");
                int sex = row.getInt("sex");

                System.out.println("id="+id+" name="+name+" age="+age+" sex="+sex);
            }
        }
    }

    @Test
    public void updateData() throws KuduException {
        KuduSession kuduSession = kuduClient.newSession();
        kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_SYNC);

        KuduTable kuduTable = kuduClient.openTable(this.kuduTable);
        Update update = kuduTable.newUpdate();
        PartialRow row = update.getRow();


        row.addInt("id", 100);
        row.addString("name", "Peter");
        row.addInt("age", 50);
        row.addInt("sex", 1);

        kuduSession.apply(update);
    }

    @Test
    public void upSertData() throws KuduException {
        KuduSession kuduSession = kuduClient.newSession();
        kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_SYNC);

        KuduTable kuduTable = kuduClient.openTable(this.kuduTable);

        Upsert upsert = kuduTable.newUpsert();
        PartialRow row = upsert.getRow();

        row.addInt("id", 100);
        row.addString("name", "Peter");
        row.addInt("age", 50);
        row.addInt("sex", 1);

        kuduSession.apply(upsert);
    }

    @Test
    public void deleteData() throws KuduException {
        KuduSession kuduSession = kuduClient.newSession();
        kuduSession.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_SYNC);

        KuduTable kuduTable = kuduClient.openTable(this.kuduTable);
        Delete delete = kuduTable.newDelete();
        PartialRow row = delete.getRow();
        row.addInt("id", 100);

        kuduSession.apply(delete);
    }

    @Test
    public void deleteTable() throws KuduException {
        String tableName = "spark_kudu_student";
        tableName = kuduTable;
        if (kuduClient.tableExists(tableName)) {
            kuduClient.deleteTable(tableName);
        }
    }

    /**
     * 创建表
     */
    @Test
    public void testPartitionsHash() throws KuduException {
        //判断表是否存在，不存在就构建
        if (!kuduClient.tableExists("t_hash_partition")) {

            //构建创建表的schema信息------就是表的字段和类型, id为表的主键
            ArrayList<ColumnSchema> columnSchemas = new ArrayList<>();
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("id", Type.INT32).key(true).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("name", Type.STRING).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("age", Type.INT32).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("sex", Type.INT32).build());
            Schema schema = new Schema(columnSchemas);

            //指定创建表的option属性
            CreateTableOptions tableOptions = new CreateTableOptions();

            //指定表的分区规则 采用hash分区 根据id哈希到指定的6个部分中
            ArrayList<String> partitionList = new ArrayList<>();
            partitionList.add("id");
            tableOptions.addHashPartitions(partitionList, 6); //id.hashcode % 6

            //如果表不存在，进行创建表的操作(表名 表schema信息 表属性信息)
            kuduClient.createTable("t_hash_partition", schema, tableOptions);
        }
    }

    @Test
    public void testPartitionsRange() throws KuduException {
        //判断表是否存在，不存在就构建
        if (!kuduClient.tableExists("t_range_partition")) {

            //构建创建表的schema信息------就是表的字段和类型, id为表的主键
            ArrayList<ColumnSchema> columnSchemas = new ArrayList<>();
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("id", Type.INT32).key(true).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("name", Type.STRING).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("age", Type.INT32).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("sex", Type.INT32).build());
            Schema schema = new Schema(columnSchemas);

            //指定创建表的option属性
            CreateTableOptions tableOptions = new CreateTableOptions();

            ArrayList<String> partitionList = new ArrayList<>();
            partitionList.add("id");

            //指定range分区的字段
            tableOptions.setRangePartitionColumns(partitionList);
            //指定分区策略
            /**
             * 0  <= id < 10
             * 10 <= id < 20
             * 20 <= id < 30
             * 30 <= id < 40
             * 40 <= id < 50
             */
            int count = 0;
            for (int i = 0; i < 5; i++) {
                //指定range的下界
                PartialRow lower = schema.newPartialRow();
                lower.addInt("id", count);
                count+=10;

                //指定range的上界
                PartialRow upper = schema.newPartialRow();
                upper.addInt("id", count);
                tableOptions.addRangePartition(lower, upper);
            }

            //如果表不存在，进行创建表的操作(表名 表schema信息 表属性信息)
            kuduClient.createTable("t_range_partition", schema, tableOptions);
        }
    }


    @Test
    public void testPartitionsMulti() throws KuduException {
        //判断表是否存在，不存在就构建
        if (!kuduClient.tableExists("t_multi_partition")) {

            //构建创建表的schema信息------就是表的字段和类型, id为表的主键
            ArrayList<ColumnSchema> columnSchemas = new ArrayList<>();
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("id", Type.INT32).key(true).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("name", Type.STRING).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("age", Type.INT32).build());
            columnSchemas.add(new ColumnSchema.ColumnSchemaBuilder("sex", Type.INT32).build());
            Schema schema = new Schema(columnSchemas);

            //指定创建表的option属性
            CreateTableOptions tableOptions = new CreateTableOptions();

            ArrayList<String> partitionList = new ArrayList<>();
            partitionList.add("id");

            tableOptions.addHashPartitions(partitionList, 3);

            //指定range分区的字段
            tableOptions.setRangePartitionColumns(partitionList);
            //指定分区策略
            /**
             * 0  <= id < 10
             * 10 <= id < 20
             * 20 <= id < 30
             * 30 <= id < 40
             * 40 <= id < 50
             */
            int count = 0;
            for (int i = 0; i < 5; i++) {
                //指定range的下界
                PartialRow lower = schema.newPartialRow();
                lower.addInt("id", count);
                count+=10;

                //指定range的上界
                PartialRow upper = schema.newPartialRow();
                upper.addInt("id", count);
                tableOptions.addRangePartition(lower, upper);
            }

            //如果表不存在，进行创建表的操作(表名 表schema信息 表属性信息)
            kuduClient.createTable("t_multi_partition", schema, tableOptions);
        }
    }
}
